from model.BaseModel import BaseModel
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import tushare as ts
import warnings

from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM, SimpleRNN
from keras.layers import Dropout
import keras
from sklearn.preprocessing import MinMaxScaler

from io import BytesIO
from PIL import Image
import base64


class LSTMModel(BaseModel):

    def __init__(self):
        self.endTime = None
        self.startTime = None
        self.code = None
        self.imgUrls = None
        self.y_test = None
        self.y_train = None
        self.x_test = None
        self.x_train = None
        self.regressor = None
        self.all_data = None
        self.history = None
        self.websocket = None

    def setWS(self, ws):
        self.websocket = ws
        pass

    def setCode(self, code):
        self.code = str(code)
        pass

    def setTime(self, start, end):
        self.startTime = start
        self.endTime = end
        pass

    async def view(self):
        print(self.code)
        print(self.startTime)
        print(self.endTime)
        data = ts.get_k_data(self.code, start=self.startTime, end=self.endTime)
        d = []
        for index, row in data.iterrows():
            print(row['date'], row['open'], row['close'], row['high'], row['low'], row['volume'])
            d.append({
                'date': row['date'],
                'open': row['open'],
                'close': row['close'],
                'high': row['high'],
                'low': row['low'],
                'volume': row['volume'],
            })

        await self.websocket.send(self.result("data", "lstm-view-data", "200", "模型训练开始", d))
        pass

    async def train(self):
        await self.websocket.send(self.result("message", "lstm-train", "200", "模型训练开始"))
        # In[6]:

        # Importing the training set读入数据集
        # 下载数据接口介绍https://mp.weixin.qq.com/s/XoyACntxEXX3ZEqvECUbBg
        data = ts.get_k_data(self.code, start=self.startTime, end=self.endTime)  # 通过tushare的接口获取浦发银行的历史数据
        print(data.shape)

        # In[7]:

        # 显示数据集的前几行，查看样式
        print(data.head())
        self.all_data = data.iloc[:, 1:6]
        print(self.all_data.head())

        # # Feature Scaling特征归一化
        #

        # In[9]:

        sc = MinMaxScaler(feature_range=(0, 1), )
        all_data_scaled = sc.fit_transform(self.all_data)
        print(all_data_scaled)
        print('训练数据长度是:', len(all_data_scaled))

        features = []
        labels = []
        for i in range(60, len(all_data_scaled)):
            features.append(all_data_scaled[i - 60:i, ])
            labels.append(all_data_scaled[i, 1])
        features, labels = np.array(features), np.array(labels)
        features = np.reshape(features, (features.shape[0], features.shape[1], -1))
        self.x_train, self.x_test, self.y_train, self.y_test = features[:1600], features[1600:], labels[:1600], labels[
                                                                                                                1600:]
        print('shape of x_train:', self.x_train.shape)
        print('shape of x_test:', self.x_test.shape)
        print('shape of y_train:', self.y_train.shape)
        print('shape of y_test:', self.y_test.shape)

        warnings.filterwarnings("ignore")

        # In[12]:

        # Initialising the LSTM
        self.regressor = Sequential()
        # Adding the first LSTM layer and some Dropout regularisation

        # LSTM(100, dropout=0.2, recurrent_dropout=0.2)
        self.regressor.add(LSTM(units=50, return_sequences=True, input_shape=(self.x_train.shape[1], 5)))
        self.regressor.add(Dropout(0.2))

        # In[13]:

        # Adding a second LSTM layer and some Dropout regularisation
        self.regressor.add(LSTM(units=50, return_sequences=True))
        self.regressor.add(Dropout(0.2))

        # In[14]:

        # Adding a third LSTM layer and some Dropout regularisation
        self.regressor.add(LSTM(units=50, return_sequences=True))
        self.regressor.add(Dropout(0.2))

        # In[15]:

        # Adding a fourth LSTM layer and some Dropout regularisation
        self.regressor.add(LSTM(units=50))
        self.regressor.add(Dropout(0.2))

        from sklearn.model_selection import train_test_split
        # Adding the output layer
        self.regressor.add(Dense(units=1))
        # Compiling the LTSM
        self.regressor.compile(optimizer='adam', loss='mean_squared_error')

        # Fitting the LSTM to the Training set
        class LossHistory(keras.callbacks.Callback):
            def on_train_begin(self, logs={}):
                self.losses = {'batch': [], 'epoch': []}
                self.accuracy = {'batch': [], 'epoch': []}
                self.val_loss = {'batch': [], 'epoch': []}
                self.val_acc = {'batch': [], 'epoch': []}

            def on_batch_end(self, batch, logs={}):
                self.losses['batch'].append(logs.get('loss'))
                self.accuracy['batch'].append(logs.get('acc'))
                self.val_loss['batch'].append(logs.get('val_loss'))
                self.val_acc['batch'].append(logs.get('val_acc'))

            def on_epoch_end(self, batch, logs={}):
                self.losses['epoch'].append(logs.get('loss'))
                self.accuracy['epoch'].append(logs.get('acc'))
                self.val_loss['epoch'].append(logs.get('val_loss'))
                self.val_acc['epoch'].append(logs.get('val_acc'))

            def loss_plot(self, loss_type):
                iters = range(len(self.losses[loss_type]))
                plt.figure()
                # loss
                plt.plot(iters, self.losses[loss_type], 'g', label='train loss')
                if loss_type == 'epoch':
                    # val_loss
                    plt.plot(iters, self.val_loss[loss_type], 'k', label='val loss')
                plt.grid(True)
                plt.xlabel(loss_type)
                plt.ylabel('acc-loss')
                plt.legend(loc="upper right")
                plt.show()

        self.history = self.regressor.fit(self.x_train, self.y_train, epochs=5, batch_size=32, validation_data=(
            self.x_test, self.y_test))
        await self.websocket.send(self.result("message", "lstm-train", "200", "模型训练结束"))
        pass

    async def draw(self):
        self.imgUrls = []
        # In[17]:
        history = self.history
        # 画损失曲线图
        loss = history.history['loss']
        val_loss = history.history['val_loss']
        epochs = range(1, len(loss) + 1)
        plt.title('Loss curve')
        plt.plot(epochs, loss, 'red', label='Training loss')
        plt.plot(epochs, val_loss, 'blue', label='Validation loss')
        # 生成第一张图表
        buffer = BytesIO()
        plt.savefig(buffer, format='png')
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        self.imgUrls.append(img_base64)
        plt.legend()
        plt.show()

        # In[16]:

        sc_one = MinMaxScaler(feature_range=(0, 1))
        sc_one.fit_transform(self.all_data.iloc[:, 1:2])
        predicted_stock_train = self.regressor.predict(self.x_train)
        predicted_stock_train = sc_one.inverse_transform(predicted_stock_train)
        predicted_stock_test = self.regressor.predict(self.x_test)
        predicted_stock_test = sc_one.inverse_transform(predicted_stock_test)
        real_price_train = sc_one.inverse_transform(np.reshape(self.y_train, (-1, 1)))
        real_price_test = sc_one.inverse_transform(np.reshape(self.y_test, (-1, 1)))

        # In[17]:

        # Visualising the train results
        plt.plot(real_price_train, color='red', label='Real Stock Price')
        plt.plot(predicted_stock_train, color='blue', label='Predicted TAT Stock Price')
        plt.title('train Stock Price Prediction')
        plt.xlabel('Time')
        plt.ylabel('Stock Price')
        # 生成第二张图表
        buffer = BytesIO()
        plt.savefig(buffer, format='png')
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        self.imgUrls.append(img_base64)
        plt.legend()
        plt.show()

        # In[18]:

        # Visualising the test results
        plt.plot(real_price_test, color='red', label='Real Stock Price')
        plt.plot(predicted_stock_test, color='blue', label='Predicted TAT Stock Price')
        plt.title('test Stock Price Prediction')
        plt.xlabel('Time')
        plt.ylabel('Stock Price')
        # 生成第三张图表
        buffer = BytesIO()
        plt.savefig(buffer, format='png')
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        self.imgUrls.append(img_base64)
        plt.legend()
        plt.show()

        # In[19]:

        from sklearn.metrics import mean_squared_error  # 均方误差
        from sklearn.metrics import mean_absolute_error  # 平方绝对误差
        mse_score = mean_squared_error(real_price_test, predicted_stock_test)
        mae_score = mean_absolute_error(real_price_test, predicted_stock_test)
        print('测试集的均方误差是:', mse_score)
        print('测试集的平方绝对误差是:', mae_score)
        # 返回测试数据
        await  self.websocket.send(self.result("data", "lstm-draw", "200", "success", {
            "imgUrls": self.imgUrls,
            "mse_score": mse_score,
            "mae_score": mae_score
        }))
        pass
